package com.cxs.distributedencrypt.util;

import ohos.security.keystore.KeyGenAlgorithmParaSpec;
import ohos.security.keystore.KeyStoreConstants;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.security.cert.CertificateException;
import java.util.Base64;
import java.util.Objects;
import java.util.Optional;

public class SecurityUtils {

    private static final String KEY_STORE = "HarmonyKeyStore";
    private static final String KEY_PAIR_ALIAS = "HarmonyKeyPair";
    private static final String ENCRYPT_ALGORITHM = "RSA/ECB/OAEPWithSHA-256AndMGF1Padding";
    private static final String PROVIDER = "HarmonyKeyStoreBCWorkaround";
    private static final String SIGNATURE_PADDING = "PSS";

    private SecurityUtils() {}

    /**
     * 加密
     * @param cleartext 明文
     * @return
     */
    public static String encrypt(String cleartext) {
        try {
            Optional<KeyPair> keyPairOptional = getSecKey();
            if(keyPairOptional.isPresent()) {
                KeyPair keyPair = keyPairOptional.get();
                Cipher cipher = Cipher.getInstance(ENCRYPT_ALGORITHM, PROVIDER);
                cipher.init(Cipher.ENCRYPT_MODE, keyPair.getPublic());
                byte[] bytes = cipher.doFinal(cleartext.getBytes(StandardCharsets.UTF_8));
                String result = Base64.getEncoder().encodeToString(bytes);
                return result;
            }
        } catch (NoSuchAlgorithmException | NoSuchPaddingException | NoSuchProviderException
                | InvalidKeyException | BadPaddingException | IllegalBlockSizeException e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 解密
     * @param ciphertext 密文
     * @return
     */
    public static String decrypt(String ciphertext) {
        if (Objects.isNull(ciphertext)) {
            return null;
        } else {
            try {
                KeyStore keyStore = KeyStore.getInstance(KEY_STORE);
                keyStore.load(null);
                Key key = keyStore.getKey(KEY_PAIR_ALIAS, null);
                PrivateKey privateKey;
                if (key instanceof PrivateKey) {
                    privateKey = (PrivateKey) key;
                } else {
                    return null;
                }
                byte[] bytes = Base64.getDecoder().decode(ciphertext);
                Cipher cipher = Cipher.getInstance(ENCRYPT_ALGORITHM, PROVIDER);
                cipher.init(Cipher.DECRYPT_MODE, privateKey);
                String result = new String(cipher.doFinal(bytes)).trim();
                return result;
            } catch (KeyStoreException | CertificateException | IOException | NoSuchAlgorithmException
                    | UnrecoverableKeyException | NoSuchPaddingException | NoSuchProviderException
                    | InvalidKeyException | BadPaddingException | IllegalBlockSizeException e) {
                e.printStackTrace();
            }
        }
        return null;
    }

    private static Optional<KeyPair> getSecKey() {
        try {
            KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(KeyStoreConstants.SEC_KEY_ALGORITHM_RSA,
                    KEY_STORE);
            KeyGenAlgorithmParaSpec.Builder builder = new KeyGenAlgorithmParaSpec.Builder(
                    KEY_PAIR_ALIAS).setSecKeyUsagePurposes(
                    KeyStoreConstants.PURPOSE_CAN_ENCRYPT | KeyStoreConstants.PURPOSE_CAN_DECRYPT
                            | KeyStoreConstants.PURPOSE_CAN_SIGN | KeyStoreConstants.PURPOSE_CAN_VERIFY)
                    .addSecKeyCryptoAttr(KeyStoreConstants.CRYPTO_PARAMETER_BLOCK_MODE,
                            KeyStoreConstants.SEC_BLOCK_MODE_ECB)
                    .addSecKeyCryptoAttr(KeyStoreConstants.CRYPTO_PARAMETER_ENCRYPT_PADDING,
                            KeyStoreConstants.OPTIMAL_ASYMMETRIC_ENCRYPTION_PADDING)
                    .addSecKeyCryptoAttr(KeyStoreConstants.CRYPTO_PARAMETER_DIGEST,
                            KeyStoreConstants.DIGEST_ALGORITHM_SHA256)
                    .addSecKeyCryptoAttr(KeyStoreConstants.CRYPTO_PARAMETER_SIGNATURE_PADDING, SIGNATURE_PADDING);
            KeyGenAlgorithmParaSpec keyGenAlgorithmParaSpec = builder.createKeyGenAlgorithmParaSpec();
            keyPairGenerator.initialize(keyGenAlgorithmParaSpec);
            return Optional.of(keyPairGenerator.generateKeyPair());
        } catch (NoSuchProviderException | NoSuchAlgorithmException | InvalidAlgorithmParameterException e) {
            return Optional.empty();
        }
    }
}
